#!/usr/bin/env python
import subprocess as sp
import os
import sys
import glob
import pandas as pd
import psutil
import re
import shutil
from itertools import islice
import argparse
import coordmagic as cm
from matplotlib.colors import LinearSegmentedColormap
import matplotlib.pyplot as plt
import numpy as np
'''todo list
do not modify log file'''


'''version 1.1  20221003
use Multiwfn to generate figures in batch
developed by zhongcheng@whu.edu.cn. QQ:32598827
please contact me if you have any suggestions or find any bugs
'''
class extractInfo:
    '''detect homo and lumo in wfn file (fchk,log,molden)
    and convert orbital index to orbital label or
    convert orbital label to orbital index'''

    def __init__(self, wfnfile):
        self.wfnfile = os.path.abspath(wfnfile)
        try:
            self.fileobj = open(wfnfile, 'r')
            filename = os.path.basename(wfnfile)
            self.basename, self.ext = os.path.splitext(filename)
        except TypeError:
            sys.exit(
                'Error! The input {:s} is not a file' .format(str(wfnfile)))

    def read_orb(self):
        self.openshell = 0
        readorb = {'.fchk': self._orb_fchk,
                   '.log': self._orb_log,
                   '.molden': self._orb_molden}
        self.homo_idx = readorb[self.ext]()
        self.total_orb = self.homo_idx[1] - self.homo_idx[0]

    def read_exc(self):
        '''es_states is a list of dict'''
        readexc = {'.log': self._exc_log,
                   '.out': self._exc_out,
                   }
        es_states = readexc[self.ext]()
        es_no = []
        es_lbl = []
        no_s = 0
        no_t = 0
        for e in es_states:
            if e['Mult'].startswith('S'):
                no_s = no_s + 1
                elbl = 'S'+str(no_s)
            elif e['Mult'].startswith('T'):
                no_t = no_t + 1
                elbl = 'T'+str(no_t)
            else:
                elbl = 'E'+str(e['sn'])
            es_no.append(e['sn'])
            es_lbl.append(elbl)
        self.es_no = es_no
        self.es_lbl = es_lbl

    def _exc_log(self):
        sep_line = [0]
        for line_no, line in enumerate(self.fileobj):
            if 'SCF Done' in line and line_no > 0:
                sep_line.append(line_no)
        sep_line.append(line_no)

        if len(sep_line) > 3:
            print('Optimize job detected, replace current gout file with its last part.')
            self.fileobj = open(self.wfnfile, 'r')
            first_part = list(islice(self.fileobj, sep_line[1]))
            self.fileobj = open(self.wfnfile, 'r')
            last_part = list(islice(self.fileobj, sep_line[-2], sep_line[-1]))
            self.fileobj.close()
            shutil.move(self.wfnfile, self.wfnfile+'.bak')
            newlog = open(self.wfnfile, 'w')
            for line in first_part+last_part:
                if 'GradGrad' not in line:
                    newlog.write(line)
            newlog.close()
        self.fileobj = open(self.wfnfile, 'r')
        excited_states = []
        for l in self.fileobj:
            if 'Excited State' in l:
                es = {'sn': int(l.split()[2].strip(':')),
                      'Mult': l.split()[3].split('-')[0],
                      }
                excited_states.append(es)
        return excited_states

    def _exc_out(self):
        pass

    def _orb_fchk(self):
        for l in self.fileobj:
            if 'Beta Orbital Energies' in l:
                self.openshell = 1
            if 'Number of alpha electrons' in l:
                alpha_e = l.split()[5]
            if 'Number of beta electrons' in l:
                beta_e = l.split()[5]
            if 'Number of independent functions' in l:
                total_orb = l.split()[5]
            if 'MO coefficients' in l:
                break
        beta_homo = int(total_orb)+int(beta_e)
        return [int(alpha_e), beta_homo]

    def _orb_log(self):
        for l in self.fileobj:
            if 'Orbital symmetries:' in l:
                line = self.fileobj.next()
                if 'Alpha Orbitals' in line:
                    self.openshell = 1
                break
            if 'alpha electrons' in l:
                alpha_e = l.split()[0]
                beta_e = l.split()[3]
            if 'primitive gaussians' in l:
                total_orb = l.split()[0]
        beta_homo = str(int(total_orb)+int(beta_e))
        return [int(alpha_e), beta_homo]

    def _orb_molden(self):
        pass


class parseOrb:
    def __init__(self, info_obj):
        self.openshell = info_obj.openshell
        self.homo_idx = info_obj.homo_idx
        self.total_orb = info_obj.total_orb

    def parse_orb(self, orb_str):
        '''convert user input orbital string to a list of orbital label and index'''
        orb_label = []
        for o in orb_str.split(','):
            if '-' in o:
                o1 = o.split('-')[0]
                o2 = o.split('-')[1]
                if 'a' in o1 and 'b' in o2:
                    print(
                        'Error!!! Orbtital format {:s}-{:s} is not supported'.format(o1, o2))
                    sys.exit()
                i1 = self.lbl2idx([o1])[0]
                i2 = self.lbl2idx([o2])[0]
                if i1 > i2:
                    i1, i2 = i2, i1
                orb_label = orb_label + self.idx2lbl(list(range(i1, i2+1)))
                try:
                    i1 = self.lbl2idx([o1])[1]
                    i2 = self.lbl2idx([o2])[1]
                    if i1 > i2:
                        i1, i2 = i2, i1
                    orb_label = orb_label + self.idx2lbl(list(range(i1, i2+1)))
                except IndexError:
                    pass
            else:
                orb = self.idx2lbl(self.lbl2idx([o]))
                orb_label = orb_label + orb
        self.orb_label = orb_label
        self.orb_index = self.lbl2idx(orb_label)
        self.filter_orb()

    def idx2lbl(self, idx_list):
        # convert orbital index list to orbital label list
        orb_label = []
        for i in idx_list:
            if i < self.total_orb:
                i = i - self.homo_idx[0]
                if i <= 0:
                    ol = 'H'+str(i*-1)+'a'
                else:
                    ol = 'L'+str(i-1)+'a'
            else:
                i = i - self.homo_idx[1]
                if i <= 0:
                    ol = 'H'+str(i*-1)+'b'
                else:
                    ol = 'L'+str(i-1)+'b'
            orb_label.append(ol)
        return orb_label

    def lbl2idx(self, label_list):
        '''convert orbital label to orbital index'''
        idx_list = []
        for ol in label_list:
            ol = ol.lower().replace('omo', '').replace('umo', '')
            try:
                i = int(ol.strip('hlab'))
            except ValueError:
                i = 0
            if ol.startswith('h'):
                idx_a = self.homo_idx[0] - abs(i)
                idx_b = self.homo_idx[1] - abs(i)
            elif ol.startswith('l'):
                idx_a = self.homo_idx[0] + 1 + abs(i)
                idx_b = self.homo_idx[1] + 1 + abs(i)
            else:
                try:
                    index = [int(ol)]
                except ValueError:
                    print('Error! Wrong orbital label format: {:s}'.format(ol))
            if not ol.isdigit():
                if ol.endswith('b'):
                    index = [idx_b]
                elif ol.endswith('a'):
                    index = [idx_a]
                else:
                    index = [idx_a, idx_b]
            idx_list = idx_list + index
        return idx_list

    def filter_orb(self):
        '''filt beta orbital label and index if self.open_shell=0'''
        if self.openshell == 0:
            self.orb_label = [ol.replace('a', '')
                              for ol in self.orb_label if 'b' not in ol]
            self.orb_index = [
                idx for idx in self.orb_index if idx < self.total_orb]
        self.orb_label = [re.sub('([HL])0', r'\1', ol)
                          for ol in self.orb_label]


class prepareM:
    def __init__(self, np=None, mem=None):
        if np:
            self.np = int(np)
        else:
            self.np = None
        if mem:
            if mem.lower().endswith('gb'):
                self.mem = float(mem.lower().strip('gb'))*1024*1024*1024
            elif mem.lower().endswith('mb'):
                self.mem = float(mem.lower().strip('mb'))*1024*1024
            elif mem.lower().endswith('kb'):
                self.mem = float(mem.lower().strip('kb'))*1024
            else:
                self.mem = float(mem)
        else:
            self.mem = None
        self.check_env()
        self.check_Mwfn_path()
        self.check_gau_path()
        self.modify_Mwfn_ini()

    def check_env(self):
        '''detect avaiable memory, physical cpu cores and Multiwfn path'''
        self.settings = {}
        self.mem_avail = float(psutil.virtual_memory()._asdict()['available'])
        print('{:.1f} GB memory avaialbe'.format(
            float(self.mem_avail)/1024/1024/1024))
        self.np_avail = int(psutil.cpu_count(logical=False))
        print('{:d} physical cores detected'.format(self.np_avail))
        if not self.np:
            self.np = self.np_avail
        elif self.np > self.np_avail:
            self.np = self.np_avail
        if not self.mem:
            self.mem = self.mem_avail
        elif self.mem > self.mem_avail:
            self.mem = self.mem_avail
        self.stacksize = int(self.mem*0.8/self.np)
        if self.stacksize > 2147483647:
            self.stacksize = 2147483647
        print('Use {:d} cores and ompstacksize is set to {:.1f} GB'.format(
            self.np, self.stacksize/1024/1024/1024))
        os.environ["KMP_STACKSIZE"] = str(self.stacksize)
        self.settings['ompstacksize'] = str(self.stacksize)
        self.settings['nthreads'] = str(self.np)

    def check_Mwfn_path(self):
        # check Multiwfn path
        path = os.getenv('PATH')
        Mdir = None
        for p in path.split(os.path.pathsep):
            p1 = os.path.join(p, 'Multiwfn.exe')
            p2 = os.path.join(p, 'Multiwfn')
            if os.path.exists(p1) and os.access(p1, os.X_OK):
                Mdir = p
                break
            elif os.path.exists(p2) and os.access(p2, os.X_OK):
                Mdir = p
                break
        if not Mdir:
            print('Error!!! Multiwfn not found in PATH, exit now!')
            os._exit(0)
        else:
            print('Multiwfn found in {:s}'.format(Mdir))
        self.Mdir = Mdir
    
    def _check_path(self,com_name):
        path = os.getenv('PATH')
        for p in path.split(os.path.pathsep):
            p1 = os.path.join(p, com_name+'.exe')
            p2 = os.path.join(p, com_name)
            if os.path.exists(p1) and os.access(p1, os.X_OK):
                return p1
            elif os.path.exists(p2) and os.access(p2, os.X_OK):
                return p2
        return None

    def check_gau_path(self):
        # check gaussian path
        path = os.getenv('PATH')
        Gdir = None
        gau = None
        formchk = self._check_path('formchk')
        cubegen = self._check_path('cubegen')
        gaupath = self._check_path('g16')
        if not gaupath:
            gaupath = self._check_path('g09')
        if gaupath:
            print('Gaussian found: {:s}'.format(gaupath))
            self.settings['gaupath'] = gaupath
        else:
            print('Warning!!! g16 or g09 not found in PATH')
        if cubegen:
            print('cubegen found: {:s}'.format(cubegen))
            self.settings['cubegenpath'] = cubegen
        else:
            print('Warning! cubegenpath not set becasue not exist or excutable'\
                  '(used to generate ESP cube)')
        if formchk:
            print('formchk found: {:s}'.format(formchk))
            self.settings['formchkpath'] = formchk
        else:
            print('Warning! formchkpath not set becasue not exist or excutable'\
                  '(only needed if u didnot formchk youself)')

    def modify_Mwfn_ini(self, **kwargs):
        '''modify settings.ini by inp_dict'''
        self.settings.update(kwargs)
        inifile = os.path.join(self.Mdir, 'settings.ini')
        if os.path.exists(inifile):
            try:
                shutil.copyfile(inifile, 'settings.ini')
            except SameFileError:
                pass
            f = open('settings.ini', 'r+')
            f_content = f.read()
            for k, v in self.settings.items():
                v = v.replace('\\', '\\\\')
                if 'path' in k:
                    #f_content = re.sub(k+'.*//', r'{}= "{}"'.format(k,v), f_content)
                    f_content = re.sub(
                        k+'.*/?/?', '{}= "{}" //'.format(k, v), f_content)
                else:
                    f_content = re.sub(
                        k+'.*/?/?', r'{}= {} //'.format(k, v), f_content)
            f.seek(0)
            f.truncate()
            f.write(f_content)
            f.close()
            cwd = os.getcwd()
            os.environ["Multiwfnpath"] = cwd
            print('settings.ini is modified and copy to {:s}'.format(cwd))
            print('Set Multiwfnpath to {:s}'.format(cwd))
        else:
            print(
                'Warning! settings.ini is not found in {:s}'.format(self.Mdir))


class prepareFile:
    def __init__(self, basename, pm):
        '''check and prepare all the necessary file
        and pass this file object to Multiwfn class
        basename is the basename of a file without extenstion
        pm is prepareM object that contains path information
        of formchk'''
        self.pm = pm
        self.basename = basename
        if os.path.exists(basename+'.log'):
            self.logfile = basename+'.log'
        elif os.path.exists(basename+'.out'):
            self.logfile = basename+'.out'
        else:
            self.logfile = ""

        if os.path.exists(basename+'.fchk'):
            self.fchkfile = basename+'.fchk'
        elif os.path.exists(basename+'.chk'):
            self._formchk(basename+'.chk')
        else:
            sys.exit('Error! {:s}.fchk or {:s}.chk not available'
                     .format(basename, basename))

    def _formchk(self, chkfile):
        if "formchkpath" in self.pm.settings:
            fchkpath = self.pm.settings["formchkpath"]
            print("formchk {:s} using {:s}".format(chkfile, fchkpath))
            p = sp.Popen([fchkpath, self.basename+'.chk'])
            p.communicate()
            self.fchkfile = self.basename+'.fchk'
        else:
            sys.exit("Error ! formchk command not found. Please formchk the chk file youself")

class Mout_Render:
    '''render / convert various figure type based on various Multiwfn output,
       all the render parameters are associate with profile'''
    def __init__(self,m,dpi=300):
        ''' m is the Multiwfn object'''
        self.basename=m.basename
        self.dpi=int(dpi)
        self.m=m
        self.p2c = {
                'nci':{'render':self.render_nci,'converter':self.ps2png},
                'ncip':{'render':self.render_nci,'converter':self.ps2png},
                'iri':{'render':self.render_iri,'converter':self.ps2png},
                'igm':{'render':self.render_igm,'converter':self.ps2png},
                'igmh':{'render':self.render_igm,'converter':self.ps2png}
                }

    def render(self):
        os.chdir("MFIG")
        for prof,mout in self.m.p2o.items():
            self.p2c[prof]['render'](mout)
        os.chdir('..')


    def check_imagemagick_path(self):
        path = os.getenv('PATH')
        Ccom = None
        for p in path.split(os.path.pathsep):
            p1 = os.path.join(p, 'convert.exe')
            p2 = os.path.join(p, 'convert')
            if os.path.exists(p1) and os.access(p1, os.X_OK):
                Ccom = p1
                break
            elif os.path.exists(p2) and os.access(p2, os.X_OK):
                Ccom = p2
                break
        if not Ccom:
            print('Error!!! mogrify not found in PATH, exit now!')
            os._exit(0)
        self.Ccom = Ccom

    def ps2png(self,inp_fig):
        self.check_imagemagick_path()
        out_fig = os.path.splitext(inp_fig)[0] +'.png'
        print('running ' + ' '.join([self.Ccom,'-density',str(self.dpi),'-rotate','90',inp_fig,out_fig]))
       # p = sp.Popen(' '.join([self.Ccom,'-density', str(self.dpi),'-rotate','90',inp_fig,out_fig]),
       #              shell=True,stdout=sp.PIPE, stderr=sp.PIPE)
        os.system(' '.join([self.Ccom,'-density', str(self.dpi),'-rotate','90',inp_fig,out_fig]))
       # stdout,stderr=p.communicate()
       # if stderr:
       #     print(stderr)
       # if stdout:
       #     print(stdout)

    def render_iri(self,mout):
        self.check_gnuplot_path()
        out_fig = os.path.splitext(mout)[0]+'.ps'
        inp_str = """set terminal postscript landscape enhanced color 'Helvetica' 20
set encoding iso_8859_1
set output '{:s}' 
set key 
set ylabel 'IRI (a.u.)' font "Helvetica, 20" 
set xlabel 'sign({{/Symbol-Oblique l}}_2){{/Symbol-Oblique r}} (a.u.)' font "Helvetica, 20"
set pm3d map
set palette defined (-0.04 "blue",0 "green", 0.02 "red")
set format y "%.1f"
set format x "%.2f"
set format cb "%.2f"
set border lw 2
set xrange [-0.05:0.05]  
set xtic  -0.05,0.01,0.05 nomirror rotate font "Helvetica"
set yrange [0.0:5.0] 
set ytic   0.0,0.5,5.0 nomirror font "Helvetica"
set cbtic  -0.04,0.01,0.02 nomirror font "Helvetica"
set cbrange [-0.04:0.02]
plot '{:s}' u 4:5:4 with points pointtype 31 pointsize 0.3 palette t ''""".format(out_fig,mout)
        gout = open('gnuplot_out', 'w')
        gerr = open('gnuplot_err', 'w')
        with open('IRIscatter.gnu','w') as gnuinp:
            gnuinp.write(inp_str)
        print('running ' + ' '.join([self.Gcom,'IRIscatter.gnu']))
        p = sp.Popen([self.Gcom,'IRIscatter.gnu'],
                     stdout=gout, universal_newlines=True, stderr=gerr)
        gout.close()
        gerr.close()
        return out_fig

    def render_igm(self,mout):
        cmap_bgr = LinearSegmentedColormap.from_list('bgr', ['#0000FF', '#00FF00', '#FF0000'])
        out_fig = os.path.splitext(mout)[0]+'_inter.png'
        dg, sl2r = np.loadtxt(mout, unpack=True, usecols=(0, 3))
        # column 3 is sl2r, which should not be changed. column 0, 1, 2 are delta_g_inter, delta_g_intra
        # and delta_g, respectively. The default is to use the delta_g_inter.
        # here if, by default, use delta_g_inter, the variable 'dg' stands for delta_g_inter.
        # if you change the line you used, you'd better change the variable name and y label as well.
        fig, ax = plt.subplots(figsize=(9.6, 7.2))
        sc = ax.scatter(sl2r, dg, c=sl2r, s=1, vmin=-0.05, vmax=0.05, cmap=cmap_bgr)
        ax.set_xlim(-0.05, 0.05)
        ax.set_ylim(0, 0.04)
        ax.set_xticks(np.linspace(-0.05, 0.05, 11))
        ax.set_yticks(np.linspace(0, 0.04, 5))
        ax.set_xlabel('$\\mathrm{sign}\\left(\\lambda_2\\right)\\rho$ (a.u.)')
        ax.set_ylabel(u'$\u03b4g_{inter}$ (a.u.)')
        cbar = plt.colorbar(sc)
        cbar.set_ticks(np.linspace(-0.05, 0.05, 11))
        fig.savefig(out_fig,dpi=self.dpi)
        # save intra fig
        out_fig = os.path.splitext(mout)[0]+'_intra.png'
        dg, sl2r = np.loadtxt(mout, unpack=True, usecols=(1, 3))
        fig, ax = plt.subplots(figsize=(9.6, 7.2))
        sc = ax.scatter(sl2r, dg, c=sl2r, s=1, vmin=-0.05, vmax=0.05, cmap=cmap_bgr)
        ax.set_xlim(-0.05, 0.05)
        ax.set_ylim(0, 0.08)
        ax.set_xticks(np.linspace(-0.05, 0.05, 11))
        ax.set_yticks(np.linspace(0, 0.08, 9))
        ax.set_xlabel('$\\mathrm{sign}\\left(\\lambda_2\\right)\\rho$ (a.u.)')
        ax.set_ylabel(u'$\u03b4g_{intra}$ (a.u.)')
        cbar = plt.colorbar(sc)
        cbar.set_ticks(np.linspace(-0.05, 0.05, 11))
        fig.savefig(out_fig,dpi=self.dpi)
        # save intra+inter fig
        out_fig = os.path.splitext(mout)[0]+'_alldg.png'
        dg, sl2r = np.loadtxt(mout, unpack=True, usecols=(2, 3))
        fig, ax = plt.subplots(figsize=(9.6, 7.2))
        sc = ax.scatter(sl2r, dg, c=sl2r, s=1, vmin=-0.05, vmax=0.05, cmap=cmap_bgr)
        ax.set_xlim(-0.05, 0.05)
        ax.set_ylim(0, 0.08)
        ax.set_xticks(np.linspace(-0.05, 0.05, 11))
        ax.set_yticks(np.linspace(0, 0.08, 9))
        ax.set_xlabel('$\\mathrm{sign}\\left(\\lambda_2\\right)\\rho$ (a.u.)')
        ax.set_ylabel(u'$\u03b4g$ (a.u.)')
        cbar = plt.colorbar(sc)
        cbar.set_ticks(np.linspace(-0.05, 0.05, 11))
        fig.savefig(out_fig,dpi=self.dpi)

    def render_nci(self,mout):
        self.check_gnuplot_path()
        out_fig = os.path.splitext(mout)[0]+'.ps'
        inp_str = """set terminal postscript landscape enhanced color 'Helvetica' 20
set encoding iso_8859_1
set output '{:s}'
set key
set ylabel 'RDG' font "Helvetica, 20"
set xlabel 'sign({{/Symbol-Oblique l}}_2){{/Symbol-Oblique r}} (a.u.)' font "Helvetica, 20"
set pm3d map
set palette defined (-0.04 "blue",0 "green", 0.02 "red")
set format y "%.1f"
set format x "%.2f"
set format cb "%.2f"
set border lw 2
set xtic  -0.05,0.01,0.05 nomirror rotate font "Helvetica"
set ytic   0.0,0.2,2.0 nomirror font "Helvetica"
set cbtic  -0.04,0.01,0.02 nomirror font "Helvetica"
set xrange [-0.05:0.05]
set yrange [0.0:2.0]
set cbrange [-0.04:0.02]
plot '{:s}' u 4:5:4 with points pointtype 31 pointsize 0.3 palette t ''""".format(out_fig,mout)
        gout = open('gnuplot_out', 'w')
        gerr = open('gnuplot_err', 'w')
        with open('RDGscatter.gnu','w') as gnuinp:
            gnuinp.write(inp_str)
        print('running ' + ' '.join([self.Gcom,'RDGscatter.gnu']))
        p = sp.Popen([self.Gcom,'RDGscatter.gnu'],
                     stdout=gout, universal_newlines=True, stderr=gerr)
        gout.close()
        gerr.close()
        return out_fig

class Multiwfn:
    def __init__(self, pf_obj, frag=None):
        '''basename is the name for log and chk, profile contol which job to run
        and which infor to extract'''
        self.inp = pf_obj
        self.basename=pf_obj.basename
        self.Mout = ''
        self.frag=frag
        self.p2c = {
            'nci': ['20;1;GS;2', 'output.txt',],
            'ncip': ['20;2;GS;2', 'output.txt'],
            'iri': ['20;4;GS;2', 'output.txt'],
            'igm': ['20;10;'+self._gen_igm_frag()+';GS;2', 'output.txt'],
            'igmp': ['20;10;'+self._gen_igm_frag()+';GS;2', 'output.txt'],
            'igmh': ['20;11;'+self._gen_igm_frag()+';GS;2', 'output.txt'],
            'igmhp': ['20;11;'+self._gen_igm_frag()+';GS;2', 'output.txt'],
        }
        self.p2o = {}  # type to output file name, this dict should be completed when  genout is ran

    def _gen_igm_frag(self):
        if not self.frag:
            return "1;a"
        else:
            return ";".join([str(len(self.frag.frag2idx))]+list(self.frag.frag2idx.values()))
    def profile_parser(self, figtype=''):
        '''convert profile to list of dict, each list represent a type of cube
        [{'type':'exc','index':1,'name':}, {'type':'orb','index':55,'name':'h'}]'''
        plist = []
        name_converter = {'rdg':'nci', 'rdgp':'ncip'}
        if figtype:
            for p in figtype.split(','):
                p=p.lower()
                if p in name_converter:
                    p = name_converter[p]
                prof = {'type': p, 'name': p}
                plist.append(prof)
        # print(explist)
        if len(plist) == 0:
            sys.exit("Error!!! no fig type specified")
        self.plist = plist

    def _gen_Minp(self, profile, dens='scf', grid='3'):
        comstr = self.p2c[profile['type']][0]
        self.Mfile = self.p2c[profile['type']][1]
        _, ext = os.path.splitext(self.Mfile)
        self.outfile = self.basename + '_' + profile['name'] + ext
        if dens != 'scf':
            comstr = '200;16;'+dens.upper()+';y;0;' + comstr
            self.outfile = self.basename + '_' + \
                dens.upper() + profile['name'] + ext
        comstr = comstr.replace('GS', grid)
        comstr = comstr.replace(';', '\n')
        m = open('Multiwfn_com', 'w')
        m.write(comstr)
        m.close()
        self.p2o[profile['type']] = self.outfile

    def _run_multiwfn(self):
        Min = open('Multiwfn_com', 'r')
        Mout = open('Multiwfn_out', 'w')
        Merr = open('Multiwfn_err', 'w')
        try:
            p = sp.Popen(['Multiwfn', self.basename+'.fchk'],
                         stdin=Min, stdout=sp.PIPE, universal_newlines=True, stderr=Merr)
        except FileNotFoundError:
            p = sp.Popen(['Multiwfn.exe', self.basename+'.fchk'],
                        stdin=Min, stdout=sp.PIPE, universal_newlines=True, stderr=Merr)
        for l in p.stdout:
            if 'Progress' in l:
                l = l.strip('\n')
                print(l, end='\r')
            else:
                Mout.write(l)
        print('\n')
        Min.close()
        Mout.close()
        Merr.close()

    def gen_out(self, grid_setting='3', dens=None):
        if not dens:
            dens = 'scf'
        if not grid_setting:
            grid_setting = '3'
        for prof in self.plist:
            self._gen_Minp(prof, grid=grid_setting, dens=dens)
            print('Using Multiwfn to generate {:s} ...'.format(self.outfile))
            self._run_multiwfn()
            if self.Mfile:
                shutil.move(self.Mfile, self.outfile)
            try:
                os.mkdir("MFIG")
            except FileExistsError:
                pass
            shutil.move(self.outfile, os.path.join("MFIG", self.outfile))
            print('{:s} moved to {:s}'.format(
                self.outfile, os.path.abspath("MFIG")))

class Fragment:
    def __init__(self,pf):
        st = cm.read_structure(pf.fchkfile)
        st.G.gen_mol(silent=True)
        self.frag2idx = {}
        self.st = st
        self.df = self.st.mol_df[["id","type_id","formula",'sn_range']]
    def r2l(self,rstr):
        l = []
        for i in rstr.split(','):
            if '-' in i:
                i1 = int(i.split('-')[0])
                i2 = int(i.split('-')[1])
                if i1 > i2:
                    i1, i2 = i2, i1
                l = l + list(range(i1, i2+1))
            else:
                try:
                    l = l + [int(i)]
                except ValueError:
                    pass
        return l

    def parser_str(self,frag_str):
        print(self.df.to_string(index=False))
        mol2idx = self.df[["id",'sn_range']].set_index('id').to_dict()['sn_range']
        mtype2idx = self.df.groupby(['type_id']).agg({'sn_range': lambda x: ','.join(list(x))}).to_dict()['sn_range']
        if frag_str == 'auto':
            self.frag2idx = mol2idx
        elif frag_str == 'autotype':
            self.frag2idx = mtype2idx
        elif frag_str != 'check':
            for i,s in enumerate(frag_str.split(';')):
                if s.startswith('m'):
                    ids = ','.join([mol2idx[i] for i in self.r2l(s.strip('m'))])
                    self.frag2idx[i+1] = ids
                elif s.startswith('t'):
                    ids = ','.join([mtype2idx[i] for i in self.r2l(s.strip('t'))])
                    self.frag2idx[i+1] = ids
                else:
                    self.frag2idx[i+1] = s
        if frag_str != 'check':
            print("Fragments:")
            print(str(self.frag2idx).replace(', ', '\n').replace(': ', ':\t').replace('{', '').replace('}', ''))

parser = argparse.ArgumentParser(
    description='Version 1.0: Use Multiwfn and Gaussian output file to generate various figures in batch', formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument('-m', dest='mem', help='Set the memory used by Multiwfn; e.g.: "-m 10GB"\n'
                                           'default: 80%% of the free memory', default=None)
parser.add_argument('-n', dest='np', help='Set the number of processors used by Multiwfn; e.g: "-n 6"\n'
                    'default: all the physical cores available.', default=None)
parser.add_argument('-d', dest='density',
                    help='choose the density to use; e.g.: "-d CI"; default is "scf"', default=None)
parser.add_argument('-t', dest='fig_type', help='Generate figure types, avail:\n'
                    'rdg,rdgp,iri,igm,igmp,igmh,igmhp,dos')
parser.add_argument('-r', dest='resolution', help='set the resolution of the final figures, '
                    'default is 300',default = "300")
parser.add_argument('-f', dest='fragment', help='define fragment, available options are:\n'
                    'check : check the molecules in current geometry\n'
                    'auto : each mol is a frag, range from large to small\n'
                    'autotype : each mol type is a frag, range from large to small\n'
                    '1,3-7;9-12 : atom index separated by semicolon\n'
                    'm1-3;m2,4 : mol index separated by semicolon\n'
                    't1,2;t3-4 : mol type index separated by semicolon\n', default='')
parser.add_argument('-g', dest='grid', help='Set the cube grid; Use same input as Multiwfn, and seperate lines by ";"\n'
                    'Some frequently used options are:\n'
                    '1/2/3 for Low Medium High quanlity grid, repectively;\n'
                    '4;x,y,z for the number of points or grid spacing in X,Y,Z direction\n'
                    '4;0.1 set the grid spaceing to 0.1 Bohr in X,Y,Z direction\n'
                    '8;abc.cube to use same grid settings as abc.cube\n'
                    'default is 2 which is the same as default cube quanlity of cubegen')
parser.add_argument('inputfile', nargs='+',
                    help='The input files are ABC.* The ABC.fchk should be available\n'
                    'Or ABC.chk with formchk in the PATH\n'
                    'ABC.log is requied if excited state is involved.')

args = parser.parse_args()
inputfiles = []
for i in args.inputfile:
    inputfiles = inputfiles + glob.glob(i)
basenames = []
for i in inputfiles:
    filename = os.path.basename(i)
    basename, ext = os.path.splitext(filename)
    basenames.append(basename)
basenames = list(dict.fromkeys(basenames))
pm = prepareM(np=args.np, mem=args.mem)
for inpf in basenames:
    pf = prepareFile(inpf, pm)
    print("Begein process file {:s} ...".format(pf.basename))
    if args.fragment:
        frag = Fragment(pf)
        frag.parser_str(args.fragment)
    else:
        frag = None
    if args.fragment != 'check':
        m = Multiwfn(pf,frag = frag)
        m.profile_parser(figtype=args.fig_type)
        m.gen_out(grid_setting=args.grid, dens=args.density)
        r = Mout_Render(m,dpi=args.resolution)
        r.render()
